# using voronoi to split the districts
import json
import numpy as np
from scipy.spatial import Voronoi
from shapely.geometry import Point, Polygon, MultiPolygon
from shapely.ops import unary_union
import geopandas as gpd

class ShenZhenVoronoiGenerator:
    def __init__(self, geojson_path: str):
        """
        初始化生成器
        geojson_path: 深圳市地理数据JSON文件路径
        """
        with open(geojson_path, 'r', encoding='utf-8') as f:
            self.data = json.load(f)
        
        self.streets_centers = []
        self.shenzhen_boundary = None
        
    def parse_polyline(self, polyline_str: str) -> Polygon | None:
        """
        解析高德地图的 polyline 字符串为 Shapely Polygon
        polyline 格式: "lng1,lat1;lng2,lat2;..." 多个区块用 | 分隔
        """
        if not polyline_str:
            return None
        
        # 处理多个多边形（用 | 分隔）
        polygons = []
        for poly_str in polyline_str.split('|'):
            coords = []
            for point in poly_str.split(';'):
                if point.strip():
                    lng, lat = map(float, point.split(','))
                    coords.append([lng, lat])
            
            if len(coords) >= 3:  # 至少需要3个点构成多边形
                # 确保首尾相连
                if coords[0] != coords[-1]:
                    coords.append(coords[0])
                polygons.append(Polygon(coords))
        
        # 合并多个多边形
        if len(polygons) == 1:
            return polygons[0]
        elif len(polygons) > 1:
            return unary_union(polygons)
        return None
    
    def extract_street_centers(self):
        """提取所有街道的中心点"""
        districts = self.data['districts'][0]['districts']
        
        for district in districts:
            district_name = district['name']
            streets = district.get('districts', [])
            
            for street in streets:
                center_str = street.get('center', '')
                if center_str:
                    lng, lat = map(float, center_str.split(','))
                    self.streets_centers.append({
                        'name': f"{district_name}-{street['name']}",
                        'district': district_name,
                        'street': street['name'],
                        'adcode': street['adcode'],
                        'center': [lng, lat]
                    })
        
        print(f"提取到 {len(self.streets_centers)} 个街道中心点")
    
    def get_shenzhen_boundary(self):
        """获取深圳市边界多边形"""
        polyline = self.data['districts'][0].get('polyline', '')
        self.shenzhen_boundary = self.parse_polyline(polyline)
        
        if self.shenzhen_boundary:
            print(f"深圳市边界面积: {self.shenzhen_boundary.area:.4f} 平方度")
        else:
            print("警告: 未能解析深圳市边界")
    
    def generate_voronoi_with_boundary(self):
        """
        生成受深圳市边界约束的 Voronoi 图
        """
        if not self.shenzhen_boundary:
            raise ValueError("请先调用 get_shenzhen_boundary()")
        
        if not self.streets_centers:
            raise ValueError("请先调用 extract_street_centers()")
        
        # 提取中心点坐标
        points = np.array([s['center'] for s in self.streets_centers])
        
        # 获取深圳边界的边界框，并扩大一些以确保 Voronoi 边界完整
        minx, miny, maxx, maxy = self.shenzhen_boundary.bounds
        margin = max(maxx - minx, maxy - miny) * 0.1  # 扩大10%
        
        # 添加边界框的四个角点（辅助点，用于限制无限远的 Voronoi 边界）
        boundary_points = np.array([
            [minx - margin, miny - margin],
            [maxx + margin, miny - margin],
            [maxx + margin, maxy + margin],
            [minx - margin, maxy + margin]
        ])
        
        all_points = np.vstack([points, boundary_points])
        
        # 计算 Voronoi 图
        print("正在计算 Voronoi 图...")
        vor = Voronoi(all_points)
        
        # 构建每个街道的 Voronoi 区域
        geometries = []
        
        for i, center_info in enumerate(self.streets_centers):
            region_index = vor.point_region[i]
            vertex_indices = vor.regions[region_index]
            
            # 跳过包含无限点的区域
            if -1 in vertex_indices or len(vertex_indices) == 0:
                print(f"跳过无效区域: {center_info['name']}")
                continue
            
            # 获取 Voronoi 区域的顶点
            vertices = vor.vertices[vertex_indices]
            voronoi_poly = Polygon(vertices)
            
            # 与深圳市边界求交集（关键步骤！）
            clipped_poly = voronoi_poly.intersection(self.shenzhen_boundary)
            
            # 跳过空的或无效的几何
            if clipped_poly.is_empty or not clipped_poly.is_valid:
                print(f"跳过空区域: {center_info['name']}")
                continue
            
            # 处理可能的 MultiPolygon（取面积最大的）
            if isinstance(clipped_poly, MultiPolygon):
                clipped_poly = max(clipped_poly.geoms, key=lambda p: p.area)
            
            geometries.append({
                'id': i + 1,
                'name': center_info['name'],
                'district': center_info['district'],
                'street': center_info['street'],
                'adcode': center_info['adcode'],
                'center_lng': center_info['center'][0],
                'center_lat': center_info['center'][1],
                'geometry': clipped_poly
            })
        
        print(f"成功生成 {len(geometries)} 个街道区域")
        
        # 转换为 GeoDataFrame
        gdf = gpd.GeoDataFrame(geometries, crs="EPSG:4326")
        return gdf
    
    def validate_coverage(self, gdf):
        """验证生成的区域是否覆盖了深圳市"""
        total_union = unary_union(gdf.geometry)
        coverage = total_union.intersection(self.shenzhen_boundary).area / self.shenzhen_boundary.area
        print(f"区域覆盖率: {coverage * 100:.2f}%")
        
        # 检查重叠
        overlaps = []
        for i in range(len(gdf)):
            for j in range(i + 1, len(gdf)):
                intersection = gdf.iloc[i].geometry.intersection(gdf.iloc[j].geometry)
                if not intersection.is_empty and intersection.area > 1e-8:
                    overlaps.append((i, j, intersection.area))
        
        if overlaps:
            print(f"发现 {len(overlaps)} 对重叠区域（Voronoi图理论上不应重叠）")
        else:
            print("✓ 无重叠区域")
    
    def export_geojson(self, gdf, output_path: str):
        """导出为 GeoJSON 文件"""
        gdf.to_file(output_path, driver='GeoJSON')
        print(f"已保存到: {output_path}")
    
    def export_to_postgis(self, gdf, db_url: str, table_name: str = 'communities'):
        """导出到 PostGIS 数据库"""
        from sqlalchemy import create_engine
        engine = create_engine(db_url)
        gdf.to_postgis(table_name, engine, if_exists='replace', index=False)
        print(f"已导入到 PostGIS 表: {table_name}")

if __name__ == "__main__":
    # 1. 初始化生成器
    data_path = "../data/raw/shenzhen_geo.json"
    export_path = "../data/processed/shenzhen_voronoi_communities.geojson"
    generator = ShenZhenVoronoiGenerator(data_path)

    # 2. 提取街道中心点
    generator.extract_street_centers()
    
    # 3. 获取深圳市边界
    generator.get_shenzhen_boundary()
    
    # 4. 生成 Voronoi 图（受边界约束）
    gdf = generator.generate_voronoi_with_boundary()
    
    # 5. 验证覆盖率
    generator.validate_coverage(gdf)
    
    # 6. 导出 GeoJSON
    generator.export_geojson(gdf, export_path)
    
    # 7. （可选）导出到数据库
    # generator.export_to_postgis(
    #     gdf, 
    #     'postgresql://user:pass@localhost:5432/bike_sharing'
    # )
    
    # 8. 可视化验证
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(15, 15))
    
    # 绘制深圳边界（修复：处理 MultiPolygon）
    if generator.shenzhen_boundary:
        if isinstance(generator.shenzhen_boundary, MultiPolygon):
            # 如果是多个多边形，分别绘制每个
            for poly in generator.shenzhen_boundary.geoms:
                x, y = poly.exterior.xy
                ax.plot(x, y, 'k-', linewidth=2)
            ax.plot([], [], 'k-', linewidth=2, label='Shenzhen Boundary')  # 图例
        else:
            # 单个多边形
            x, y = generator.shenzhen_boundary.exterior.xy
            ax.plot(x, y, 'k-', linewidth=2, label='Shenzhen Boundary')
    
    # 绘制 Voronoi 区域
    gdf.plot(ax=ax, column='district', alpha=0.5, edgecolor='black', 
             linewidth=0.5, legend=True, cmap='tab20')
    
    # 绘制街道中心点
    centers = [[s['center'][0], s['center'][1]] for s in generator.streets_centers]
    centers = np.array(centers)
    ax.scatter(centers[:, 0], centers[:, 1], c='red', s=20, 
               zorder=5, label='Street Centers')
    
    ax.set_title('Shenzhen Voronoi Communities (Boundary Clipped)', fontsize=16)
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    ax.legend()
    ax.set_aspect('equal')
    
    plt.tight_layout()
    plt.savefig('../data/processed/shenzhen_voronoi_map.png', dpi=300, bbox_inches='tight')
    print("可视化图片已保存: ../data/processed/shenzhen_voronoi_map.png")
    
    # 打印统计信息
    print("\n=== 统计信息 ===")
    print(f"总区域数: {len(gdf)}")
    print(f"平均区域面积: {gdf.geometry.area.mean():.6f} 平方度")
    print(f"最大区域面积: {gdf.geometry.area.max():.6f} 平方度")
    print(f"最小区域面积: {gdf.geometry.area.min():.6f} 平方度")
    print("\n各区街道数量:")
    print(gdf['district'].value_counts())